#include "ymbot_e_control/HOFA.hpp"

// 实现 CoefficientMatrix 类的方法
CoefficientMatrix::CoefficientMatrix(const std::string& urdf_filename) {
    pinocchio::urdf::buildModel(urdf_filename, model);
    data = pinocchio::Data(model);
}

void CoefficientMatrix::print_model_info() const {
    std::cout << "Model name: " << model.name << std::endl;
    std::cout << "Number of joints: " << model.njoints << std::endl;
    std::cout << "Gravity vector from URDF: " << model.gravity.linear().transpose() << std::endl;
}

void CoefficientMatrix::compute_inertia_matrix(const Eigen::VectorXd& q) {
    pinocchio::crba(model, data, q);
    data.M.triangularView<Eigen::StrictlyLower>() = data.M.transpose().triangularView<Eigen::StrictlyLower>();
    M = data.M;
    // std::cout << "Inertia matrix (MDG coefficients):" << std::endl;
    // std::cout << M << std::endl;
    // std::cout << "Dimensions of inertia matrix M: " << M.rows() << "x" << M.cols() << std::endl;
}

void CoefficientMatrix::compute_coriolis_matrix(const Eigen::VectorXd& q, const Eigen::VectorXd& v) {
    pinocchio::computeCoriolisMatrix(model, data, q, v);
    C = data.C;
    // std::cout << "Coriolis matrix:" << std::endl;
    // std::cout << C << std::endl;
    // std::cout << "Dimensions of Coriolis matrix C: " << C.rows() << "x" << C.cols() << std::endl;
}

void CoefficientMatrix::compute_gravity_vector(const Eigen::VectorXd& q) {
    pinocchio::computeGeneralizedGravity(model, data, q);
    G = data.g;
    // std::cout << "Gravity vector:" << std::endl;
    // std::cout << G.transpose() << std::endl;
    // std::cout << "Dimensions of gravity vector G: " << G.size() << std::endl;
}

void CoefficientMatrix::verify_inertia_matrix_symmetry() const {
    if (M.isApprox(M.transpose())) {
        std::cout << "Inertia matrix is symmetric." << std::endl;
    }
    else {
        std::cout << "Inertia matrix is not symmetric." << std::endl;
    }
}

// 实现 HOFA 类的方法
HOFA::HOFA(const int& n) {
    sys_dimension = n;
    uc = Eigen::VectorXd::Zero(n, 1);
    uf = Eigen::VectorXd::Zero(n, 1);
    u = Eigen::VectorXd::Zero(n, 1);
};


void HOFA::construct_controller(const Eigen::MatrixXd& F,
                                const Eigen::MatrixXd& Z,
                                const Eigen::MatrixXd& M,
                                const Eigen::MatrixXd& C,
                                const Eigen::MatrixXd& G,
                                const Eigen::VectorXd& qe,
                                const Eigen::VectorXd& dqe,
                                const Eigen::VectorXd& d_estimated) {
    unsigned int n = Z.rows(); // Z 的行数


    // 计算 V
    Eigen::MatrixXd V(2 * n, 2 * n); // V 矩阵
    Eigen::MatrixXd ZF = Z * F;
    V << Z,
        ZF; // [Z, ZF]^T 拼接

    // std::cout << "Z: " << std::endl << Z << std::endl;
    // std::cout << "ZF: " << std::endl << ZF << std::endl;
    // std::cout << "V: " << std::endl << V << std::endl;
    // std::cout << Z.rows() << "\t " << Z.cols() << std::endl;              // 30 60
    // std::cout << ZF.rows() << "\t " << ZF.cols() << std::endl;            // 30 60
    // std::cout << n << "\t" << V.rows() << "\t " << V.cols() << std::endl; // 30 60 60

    // 计算 W
    Eigen::MatrixXd B = Eigen::MatrixXd::Identity(30, 30);
    Eigen::MatrixXd W = B.inverse() * (M * Z * F * F + C * Z * F);

    // 计算 K0 和 K1
    Eigen::MatrixXd K = W * V.inverse(); // [K0, K1]

    // std::cout << "W: " << W << std::endl;
    // std::cout << "V.inverse(): " << std::endl << V.inverse() << std::endl;
    // std::cout << "K: " << K << std::endl;

    // 拆分 K0 和 K1
    Eigen::MatrixXd K0 = K.block(0, 0, 30, 30);
    Eigen::MatrixXd K1 = K.block(0, 30, 30, 30);

    // 计算控制力矩
    uc = B.inverse() * (G + d_estimated);
    uf = K0 * qe + K1 * dqe;
    u = uc + uf;
};


NonlinearDisturbanceObserver::NonlinearDisturbanceObserver(const Eigen::MatrixXd& mp) {
    // 初始化大小为 30x30 的矩阵 m_p 和 30x1 的向量 z
    m_p = mp;
    z = Eigen::VectorXd::Zero(m_p.rows());
    d_estimated = Eigen::VectorXd::Zero(m_p.rows());
    d_simple = Eigen::VectorXd::Zero(m_p.rows());
}

// 计算 z 的数值解（欧拉法）, 计算 d 的估计值
void NonlinearDisturbanceObserver::estimate_d(const Eigen::MatrixXd& M,
                                              const Eigen::MatrixXd& C,
                                              const Eigen::MatrixXd& G,
                                              const Eigen::VectorXd& ddq,
                                              const Eigen::VectorXd& u,
                                              const double& dt) {
    Eigen::VectorXd p = m_p * ddq;
    Eigen::MatrixXd L = m_p * M.inverse();

    // 计算 z 的导数, 欧拉法积分：z_new = z_old + dz_dt * dt,  M*ddq + C*dq + G = u + d
    // Eigen::VectorXd dz_dt = -L * z + L * (C * ddq + G - u - p);
    // z = z + dz_dt * dt;
    // d_estimated = z + p;

    // 计算 z 的导数, 欧拉法积分：z_new = z_old + dz_dt * dt, M*ddq + C*dq + G + d = u
    Eigen::VectorXd dz_dt = -L * z - L * (C * ddq + G - u - p);
    z = z + dz_dt * dt;
    d_estimated = z - p;

    // std::cout << "dz_dt: " << dz_dt.transpose() << std::endl;
    // std::cout << "dt: " << dt << std::endl;
    // std::cout << "z: " << z.transpose() << std::endl;
}


// 计算 d 的估计值，使用加速信号
void NonlinearDisturbanceObserver::simple_d(const Eigen::MatrixXd& M,
                                            const Eigen::MatrixXd& C,
                                            const Eigen::MatrixXd& G,
                                            const Eigen::VectorXd& dq,
                                            const Eigen::VectorXd& ddq,
                                            const Eigen::VectorXd& u,
                                            const double& dt) {
    Eigen::MatrixXd L = m_p;

    // 计算 z 的导数
    // Eigen::VectorXd dd_dt = -L * d_simple + L * (M * ddq + C * dq + G - u); // M*ddq + C*dq + G = u + d
    Eigen::VectorXd dd_dt = -L * d_simple - L * (M * ddq + C * dq + G - u); // M*ddq + C*dq + G + d = u

    // 欧拉法积分：z_new = z_old + dz_dt * dt
    d_simple = d_simple + dd_dt * dt;
}


// test
int main(int argc, char** argv) {
    ros::init(argc, argv, "HOFA_test");
    ros::NodeHandle nh;

    // std::string package_path = ros::package::getPath("ymbot_e_description");
    // std::string urdf_filename = package_path + "/urdf/ymbot_e_description.urdf";

    // CoefficientMatrix cm(urdf_filename);
    // cm.print_model_info();

    // Eigen::VectorXd qe = Eigen::VectorXd::Constant(cm.model.nq, 3.14);
    // Eigen::VectorXd dqe = Eigen::VectorXd::Constant(cm.model.nv, 0.1);
    // cm.compute_inertia_matrix(qe);
    // cm.compute_coriolis_matrix(qe, dqe);
    // cm.compute_gravity_vector(qe);
    // cm.verify_inertia_matrix_symmetry();


    // Eigen::VectorXd diag_values = Eigen::VectorXd::Constant(30 * 2, -10);
    // Eigen::MatrixXd F = diag_values.asDiagonal();
    // Eigen::MatrixXd Z = Eigen::MatrixXd::Random(30, 30 * 2);

    // // 确保 diag_F 的大小与 Z 的列数匹配
    // if (diag_values.size() != Z.cols()) {
    //     throw std::invalid_argument("diag_F 的大小必须与 Z 的列数一致。");
    // }


    // // 创建 HOFA 对象
    // Eigen::MatrixXd M = Eigen::MatrixXd::Random(30, 30);
    // Eigen::MatrixXd C = Eigen::MatrixXd::Random(30, 30);
    // Eigen::MatrixXd G = Eigen::MatrixXd::Random(30, 1);
    // HOFA hofa(30);
    // hofa.construct_controller(F, Z, M, C, G, qe, dqe);


    // // 初始化 Z 和 F
    // Eigen::MatrixXd Z(3, 3); // 假设 Z 是 3x3 矩阵
    // Z << 1, 2, 3, 4, 5, 6, 7, 8, 9;

    // Eigen::MatrixXd F(3, 3); // 假设 F 是 3x3 矩阵
    // F << 1, 0, 0, 0, 1, 0, 0, 0, 1;

    // // 计算 ZF
    // Eigen::MatrixXd ZF = Z * F; // 矩阵乘法

    // // 拼接 V = [Z, ZF]^T
    // Eigen::MatrixXd V(6, 3); // 初始化目标矩阵 V（6 行 3 列，Z 和 ZF 的行数之和）
    // V << Z, Z * F;           // 行拼接 [Z; ZF]

    // // 输出结果
    // std::cout << "Z:\n" << Z << "\n";
    // std::cout << "ZF:\n" << ZF << "\n";
    // std::cout << "V:\n" << V << "\n";


    return 0;
}
